import os
import re
import zipfile
import random
import string
import tempfile
import base64
import mimetypes
from typing import List, Optional, Tuple
from bs4 import BeautifulSoup


class ZipMdMergeService:
    """
    ZIP文件中Markdown文件合并服务
    整合了HTML转MD和文件合并功能
    """
    
    def __init__(self, upload_folder: str = "uploads", output_folder: str = "outputs"):
        self.upload_folder = upload_folder
        self.output_folder = output_folder
        self._ensure_folders()
    
    def _ensure_folders(self):
        """确保必要的文件夹存在"""
        os.makedirs(self.upload_folder, exist_ok=True)
        os.makedirs(self.output_folder, exist_ok=True)
    
    def generate_random_string(self, length: int = 8) -> str:
        """生成指定长度的随机字符串"""
        letters = string.ascii_lowercase + string.digits
        return ''.join(random.choice(letters) for _ in range(length))
    
    def extract_page_number(self, filename: str) -> int:
        """从文件名中提取页码数字"""
        match = re.search(r'page_(\d+)', filename)
        if match:
            return int(match.group(1))
        return -1
    
    def escape_markdown(self, text: str) -> str:
        """转义 Markdown 特殊字符，确保表格格式正确"""
        if not text:
            return ''
        escape_chars = {
            '|': '\\|',
            '*': '\\*',
            '_': '\\_',
            '#': '\\#',
            '>': '\\>'
        }
        for char, escaped in escape_chars.items():
            text = text.replace(char, escaped)
        return text
    
    def html_table_to_markdown(self, html: str, br_separator: str = " ") -> str:
        """将 HTML 表格转换为 Markdown 格式"""
        if not html.strip():
            return ""

        soup = BeautifulSoup(html, "html.parser")
        tables = soup.find_all("table", recursive=True)

        if not tables:
            return ""

        result = []
        for table in tables:
            all_rows = []
            for element in table.find_all(['thead', 'tbody', 'tr']):
                if element.name == 'tr':
                    all_rows.append(element)
                else:
                    all_rows.extend(element.find_all('tr'))

            if not all_rows:
                continue

            grid = [[] for _ in range(len(all_rows))]
            
            for row_idx, tr in enumerate(all_rows):
                cells = tr.find_all(['td', 'th'])
                cell_idx = 0
                for cell in cells:
                    while cell_idx < len(grid[row_idx]) and grid[row_idx][cell_idx] is not None:
                        cell_idx += 1
                    
                    text = cell.get_text(separator=br_separator).strip()
                    text = self.escape_markdown(text)
                    rowspan = int(cell.get('rowspan', 1))
                    colspan = int(cell.get('colspan', 1))
                    
                    for i in range(rowspan):
                        for j in range(colspan):
                            target_row = row_idx + i
                            target_col = cell_idx + j
                            
                            while len(grid) <= target_row:
                                grid.append([])
                            while len(grid[target_row]) <= target_col:
                                grid[target_row].append(None)
                            
                            grid[target_row][target_col] = {
                                'text': text if i == 0 and j == 0 else '',
                                'is_source': (i == 0 and j == 0)
                            }
                    
                    cell_idx += colspan
            
            max_cols = max(len(row) for row in grid) if grid else 0
            for row in grid:
                while len(row) < max_cols:
                    row.append(None)
            
            md_lines = []
            if grid:
                header_cells = [cell['text'] if cell else '' for cell in grid[0]]
                md_lines.append("| " + " | ".join(header_cells) + " |")
                md_lines.append("| " + " | ".join(['---'] * max_cols) + " |")
                
                for row in grid[1:]:
                    data_cells = [cell['text'] if cell else '' for cell in row]
                    md_lines.append("| " + " | ".join(data_cells) + " |")
            
            result.append("\n".join(md_lines))
        
        return "\n\n".join(result)
    
    def extract_and_convert_html_tables(self, md_content: str, br_separator: str = " ") -> str:
        """从 Markdown 内容中提取 HTML 表格并转换为 Markdown 表格"""
        table_pattern = re.compile(r'<table\b.*?</table>', re.DOTALL)
        
        def replace_table(match):
            html_table = match.group(0)
            return self.html_table_to_markdown(html_table, br_separator)
        
        return table_pattern.sub(replace_table, md_content)
    
    def extract_base64_images(self, md_content: str, output_dir: str) -> str:
        """
        提取Markdown内容中的base64编码图片，保存为本地文件并替换引用；
        确保所有base64图片数据（包括畸形或无效数据）从Markdown中移除
        
        Args:
            md_content: 包含base64图片的Markdown内容
            output_dir: 输出目录路径
            
        Returns:
            替换后的Markdown内容，不包含任何base64图片数据
        """
        # 创建images子目录
        images_dir = os.path.join(output_dir, "images")
        os.makedirs(images_dir, exist_ok=True)
        
        # 匹配Markdown格式的base64图片，极宽松匹配以捕获任何以data:image;base64,开头的模式
        base64_pattern = re.compile(
            r'!\[([^\]]*?)\]\(data:image(?:/[^;)]*?)?;base64,[^\)]*?(?:\)|$)',
            re.IGNORECASE | re.DOTALL
        )
        
        def replace_base64_image(match):
            alt_text = match.group(1)
            # 尝试提取图片类型和base64数据
            try:
                full_match = match.group(0)
                data_part = full_match.split('base64,')[1].rstrip(')')
                # 提取图片类型，默认为png
                image_type_match = re.search(r'data:image/([^;]+);base64', full_match, re.IGNORECASE)
                image_type = image_type_match.group(1).lower() if image_type_match else 'png'
                
                # 清理base64数据
                base64_data = re.sub(r'\s+', '', data_part)
                
                # 解码base64数据
                image_data = base64.b64decode(base64_data, validate=True)
                
                # 验证图片类型
                if image_type not in ['png', 'jpg', 'jpeg', 'gif', 'bmp']:
                    print(f"⚠️ 不支持的图片类型: {image_type}")
                    return ""
                
                # 生成唯一的文件名
                image_filename = f"image_{self.generate_random_string(8)}.{image_type}"
                image_path = os.path.join(images_dir, image_filename)
                
                # 保存图片文件
                with open(image_path, 'wb') as f:
                    f.write(image_data)
                
                print(f"✅ 成功保存图片: {image_filename} ({len(image_data)} 字节)")
                
                # 返回新的Markdown图片引用（相对路径）
                relative_path = f"images/{image_filename}"
                return f"![{alt_text}]({relative_path})"
                
            except Exception as e:
                # 如果解码失败或数据畸形，移除base64图片
                print(f"⚠️ 无法处理base64图片: {str(e)}")
                return ""
        
        # 替换所有Markdown格式的base64图片
        processed_content = base64_pattern.sub(replace_base64_image, md_content)
        
        # 匹配HTML格式的base64图片 <img src="data:image/...">
        html_base64_pattern = re.compile(
            r'<img[^>]*src=["\']data:image(?:/[^;]*?)?;base64,[^"\'>]*["\'][^>]*>',
            re.IGNORECASE | re.DOTALL
        )
        
        def replace_html_base64_image(match):
            # 尝试提取图片类型和base64数据
            try:
                full_match = match.group(0)
                data_part = full_match.split('base64,')[1].split('"')[0].split("'")[0]
                image_type_match = re.search(r'data:image/([^;]+);base64', full_match, re.IGNORECASE)
                image_type = image_type_match.group(1).lower() if image_type_match else 'png'
                
                # 清理base64数据
                base64_data = re.sub(r'\s+', '', data_part)
                
                # 解码base64数据
                image_data = base64.b64decode(base64_data, validate=True)
                
                # 验证图片类型
                if image_type not in ['png', 'jpg', 'jpeg', 'gif', 'bmp']:
                    print(f"⚠️ 不支持的HTML图片类型: {image_type}")
                    return ""
                
                # 生成唯一的文件名
                image_filename = f"image_{self.generate_random_string(8)}.{image_type}"
                image_path = os.path.join(images_dir, image_filename)
                
                # 保存图片文件
                with open(image_path, 'wb') as f:
                    f.write(image_data)
                
                print(f"✅ 成功保存HTML图片: {image_filename} ({len(image_data)} 字节)")
                
                # 返回新的Markdown图片引用（相对路径）
                relative_path = f"images/{image_filename}"
                return f"![图片]({relative_path})"
                
            except Exception as e:
                # 如果解码失败，移除base64图片
                print(f"⚠️ 无法处理HTML base64图片: {str(e)}")
                return ""
        
        # 替换所有HTML格式的base64图片
        processed_content = html_base64_pattern.sub(replace_html_base64_image, processed_content)
        
        # 多轮清理：移除任何剩余的base64数据
        # 第一轮：清理Markdown格式的base64图片（包括畸形）
        general_base64_pattern_md = re.compile(
            r'!\[.*?\]\(data:image(?:/[^;)]*?)?;base64,[^\)]*?(?:\)|$)',
            re.IGNORECASE | re.DOTALL
        )
        processed_content = general_base64_pattern_md.sub('', processed_content)
        
        # 第二轮：清理任何剩余的data:image;base64,...数据
        general_base64_pattern_raw = re.compile(
            r'data:image(?:/[^;]*?)?;base64,[A-Za-z0-9+/=\s]*',
            re.IGNORECASE | re.DOTALL
        )
        processed_content = general_base64_pattern_raw.sub('', processed_content)
        
        return processed_content
    
    def process_zip_file(self, zip_file_path: str, transfer_html: bool = True, 
                        max_files_per_group: int = 10) -> Tuple[List[str], str]:
        """
        处理ZIP文件中的MD文件并合并，生成包含MD文件和images目录的ZIP文件
        
        Args:
            zip_file_path: ZIP文件路径（必须是字符串）
            transfer_html: 是否转换HTML表格为Markdown
            max_files_per_group: 每组最大文件数
            
        Returns:
            Tuple containing:
            - List of generated markdown file paths
            - Path to the generated ZIP file containing markdown files and images
        """
        # 验证 zip_file_path 是字符串
        if not isinstance(zip_file_path, str):
            raise ValueError(f"zip_file_path 必须是字符串，实际收到: {type(zip_file_path)}")
        
        md_files_content = {}
        
        # 创建唯一的输出目录
        output_id = self.generate_random_string()
        output_dir = os.path.join(self.output_folder, f"output_{output_id}")
        os.makedirs(output_dir, exist_ok=True)
        
        try:
            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                file_list = zip_ref.namelist()
                md_files = [f for f in file_list if f.endswith('.md') and 'nohf' not in f]
                md_files.sort(key=self.extract_page_number)
                
                for md_file in md_files:
                    with zip_ref.open(md_file) as file:
                        content = file.read().decode('utf-8')
                        if transfer_html:
                            converted_content = self.extract_and_convert_html_tables(content)
                        else:
                            converted_content = content
                        
                        # 处理base64图片
                        converted_content = self.extract_base64_images(converted_content, output_dir)
                        
                        page_number = self.extract_page_number(md_file)
                        md_files_content[page_number] = converted_content
        
        except Exception as e:
            raise Exception(f"处理ZIP文件时出错: {str(e)}")
        
        # 按页码分组处理
        page_numbers = sorted(md_files_content.keys())
        groups = [page_numbers[i:i + max_files_per_group] 
                 for i in range(0, len(page_numbers), max_files_per_group)]
        
        output_files = []
        for i, group in enumerate(groups):
            start_page = min(group)
            end_page = max(group)
            
            if len(group) == 1:
                output_filename = os.path.join(output_dir, f"merged_page_{start_page}.md")
            else:
                output_filename = os.path.join(output_dir, f"merged_pages_{start_page}_to_{end_page}.md")
            
            with open(output_filename, 'w', encoding='utf-8') as f:
                for page_number in group:
                    f.write(md_files_content[page_number])
                    f.write('\n\n')
            
            output_files.append(output_filename)
        
        # 创建包含MD文件和images目录的ZIP文件
        zip_output_path = os.path.join(self.output_folder, f"output_{output_id}.zip")
        with zipfile.ZipFile(zip_output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            # 添加MD文件
            for output_file in output_files:
                arcname = os.path.relpath(output_file, self.output_folder)
                zipf.write(output_file, arcname)
            
            # 添加images目录中的所有文件
            images_dir = os.path.join(output_dir, "images")
            if os.path.exists(images_dir):
                for root, _, files in os.walk(images_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, self.output_folder)
                        zipf.write(file_path, arcname)
        
        return output_files, zip_output_path
    
    def process_directory(self, directory_path: str, transfer_html: bool = True,
                         max_files_per_group: int = 10) -> Tuple[List[str], str]:
        """
        处理目录中的MD文件并合并，生成包含MD文件和images目录的ZIP文件
        
        Args:
            directory_path: 目录路径
            transfer_html: 是否转换HTML表格为Markdown
            max_files_per_group: 每组最大文件数
            
        Returns:
            Tuple containing:
            - List of generated markdown file paths
            - Path to the generated ZIP file containing markdown files and images
        """
        # 验证 directory_path 是字符串
        if not isinstance(directory_path, str):
            raise ValueError(f"directory_path 必须是字符串，实际收到: {type(directory_path)}")
        
        md_files_content = {}
        
        # 创建唯一的输出目录
        output_id = self.generate_random_string()
        output_dir = os.path.join(self.output_folder, f"output_{output_id}")
        os.makedirs(output_dir, exist_ok=True)
        
        try:
            file_list = os.listdir(directory_path)
            md_files = [f for f in file_list if f.endswith('.md') and 'nohf' not in f]
            md_files.sort(key=self.extract_page_number)
            
            for md_file in md_files:
                file_path = os.path.join(directory_path, md_file)
                with open(file_path, 'r', encoding='utf-8') as file:
                    content = file.read()
                    if transfer_html:
                        converted_content = self.extract_and_convert_html_tables(content)
                    else:
                        converted_content = content
                    
                    # 处理base64图片
                    converted_content = self.extract_base64_images(converted_content, output_dir)
                    
                    page_number = self.extract_page_number(md_file)
                    md_files_content[page_number] = converted_content
        
        except Exception as e:
            raise Exception(f"处理目录时出错: {str(e)}")
        
        # 按页码分组处理
        page_numbers = sorted(md_files_content.keys())
        groups = [page_numbers[i:i + max_files_per_group] 
                 for i in range(0, len(page_numbers), max_files_per_group)]
        
        output_files = []
        for i, group in enumerate(groups):
            start_page = min(group)
            end_page = max(group)
            
            if len(group) == 1:
                output_filename = os.path.join(output_dir, f"merged_page_{start_page}.md")
            else:
                output_filename = os.path.join(output_dir, f"merged_pages_{start_page}_to_{end_page}.md")
            
            with open(output_filename, 'w', encoding='utf-8') as f:
                for page_number in group:
                    f.write(md_files_content[page_number])
                    f.write('\n\n')
            
            output_files.append(output_filename)
        
        # 创建包含MD文件和images目录的ZIP文件
        zip_output_path = os.path.join(self.output_folder, f"output_{output_id}.zip")
        with zipfile.ZipFile(zip_output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            # 添加MD文件
            for output_file in output_files:
                arcname = os.path.relpath(output_file, self.output_folder)
                zipf.write(output_file, arcname)
            
            # 添加images目录中的所有文件
            images_dir = os.path.join(output_dir, "images")
            if os.path.exists(images_dir):
                for root, _, files in os.walk(images_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, self.output_folder)
                        zipf.write(file_path, arcname)
        
        return output_files, zip_output_path